import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Tokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
from datasets import Dataset

class LoRAGPT2:
    def __init__(self, model_name=r"I:\models\gpt2\models--openai-community--gpt2"):
        self.base_model = GPT2LMHeadModel.from_pretrained(model_name)
        self.tokenizer = GPT2Tokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token

        # 配置LoRA
        lora_config = LoraConfig(
            r=16,
            lora_alpha=32,
            target_modules=["c_attn", "c_proj"],  # GPT-2的注意力模块
            lora_dropout=0.1,
            bias="none",
        )

        self.model = get_peft_model(self.base_model, lora_config)
        self.model.print_trainable_parameters()

    def forward(self, input_ids, attention_mask, labels=None):
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
            output_hidden_states=True,  # 重要：返回隐藏状态
            return_dict=True
        )

        return outputs


# 使用示例
lora_gpt2 = LoRAGPT2()

# 训练配置
training_args = TrainingArguments(
    output_dir="./lora-gpt2",
    per_device_train_batch_size=4,
    num_train_epochs=3,
    learning_rate=1e-4,
    logging_dir="./logs",
    save_steps=500,
    remove_unused_columns=False,
)

# 准备训练数据
data = {
    "text":[
        "今天是个好天气。",
        "我喜欢用GPT模型学习。",
        "微调技术让模型更加灵活。",
        "LoRA 技术是一种高效的微调方法。",
        "通过低秩矩阵分解减少参数量。"
    ]
}
#转换为Hugging face数据集
dataset = Dataset.from_dict(data)

def preprocess_function(examples):
    return lora_gpt2.tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=512,
    )


tokenized_dataset = dataset.map(preprocess_function, batched=True)

